package com.szml.user.config;

import com.alibaba.druid.pool.DruidDataSource;
import com.zaxxer.hikari.HikariDataSource;
import org.apache.shardingsphere.driver.api.ShardingSphereDataSourceFactory;
import org.apache.shardingsphere.infra.config.algorithm.ShardingSphereAlgorithmConfiguration;
import org.apache.shardingsphere.sharding.api.config.ShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.strategy.keygen.KeyGenerateStrategyConfiguration;
import org.apache.shardingsphere.sharding.api.config.strategy.sharding.StandardShardingStrategyConfiguration;
import org.apache.shardingsphere.sharding.api.sharding.standard.PreciseShardingValue;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.core.env.Environment;
import io.seata.rm.datasource.DataSourceProxy;
import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.*;

/**
 * @Author limin
 * @date 2025-06-21 00:20
 * @description
 */
@Configuration
public class DataSourceConfig {

    // 创建基础数据源配置
    private DruidDataSource createBaseDataSource(String url) {
        DruidDataSource ds = new DruidDataSource();
        ds.setUrl(url);
        ds.setUsername("root");
        ds.setPassword("123456");
        ds.setDriverClassName("com.mysql.cj.jdbc.Driver");
        return ds;
    }

    @Bean("ds0")
    public DataSource ds0() {
        return createBaseDataSource(
                "jdbc:mysql://localhost:3306/user_db0?useSSL=false&serverTimezone=UTC&characterEncoding=utf8&allowPublicKeyRetrieval=true"
        );
    }

    @Bean("ds1")
    public DataSource ds1() {
        return createBaseDataSource(
                "jdbc:mysql://localhost:3306/user_db1?useSSL=false&serverTimezone=UTC&characterEncoding=utf8&allowPublicKeyRetrieval=true"
        );
    }

    @Bean("ds2")
    public DataSource ds2() {
        return createBaseDataSource(
                "jdbc:mysql://localhost:3306/user_db2?useSSL=false&serverTimezone=UTC&characterEncoding=utf8&allowPublicKeyRetrieval=true"
        );
    }


    @Primary
    @Bean("dataSource")
    public DataSource shardingSphereDataSource(
            @Qualifier("ds0") DataSource ds0,
            @Qualifier("ds1") DataSource ds1,
            @Qualifier("ds2") DataSource ds2
            ) throws SQLException {

        // 1. 创建数据源映射
        Map<String, DataSource> dataSourceMap = new HashMap<>();
        dataSourceMap.put("ds0", ds0);
        dataSourceMap.put("ds1", ds1);
        dataSourceMap.put("ds2", ds2);

        // 2. 分片规则配置
        ShardingRuleConfiguration shardingRuleConfig = new ShardingRuleConfiguration();

        // 3.1 添加广播表配置（关键部分）
        shardingRuleConfig.getBroadcastTables().add("user_roles");
        shardingRuleConfig.getBroadcastTables().add("roles");

        // 3.2 用户表分片规则
        ShardingTableRuleConfiguration userTableRule = new ShardingTableRuleConfiguration(
                "users",  // 逻辑表名
                "ds${0..2}.user_${0..2}" // 实际数据节点
        );

        // 分库策略（按user_id取模）
        userTableRule.setDatabaseShardingStrategy(
                new StandardShardingStrategyConfiguration("user_id", "db-sharding"));

        // 分表策略（按user_id取模）
        userTableRule.setTableShardingStrategy(
                new StandardShardingStrategyConfiguration("user_id", "table-sharding"));

        // 主键生成策略（关键修改）
        userTableRule.setKeyGenerateStrategy(
                new KeyGenerateStrategyConfiguration("user_id", "snowflake"));

        shardingRuleConfig.getTables().add(userTableRule);

        // 3. 分片算法配置
        Properties dbShardingProps = new Properties();
        dbShardingProps.setProperty("algorithm-expression", "ds${user_id % 3}");

        Properties tableShardingProps = new Properties();
        tableShardingProps.setProperty("algorithm-expression", "user_${user_id % 3}"); // 简化分表算法

        shardingRuleConfig.getShardingAlgorithms().put("db-sharding",
                new ShardingSphereAlgorithmConfiguration("INLINE", dbShardingProps));

        shardingRuleConfig.getShardingAlgorithms().put("table-sharding",
                new ShardingSphereAlgorithmConfiguration("INLINE", tableShardingProps));

        // 4. 分布式序列配置（关键修改）
        Properties snowflakeProps = new Properties();
        snowflakeProps.setProperty("worker-id", "123");
        snowflakeProps.setProperty("max-vibration-offset", "3"); // 防止冲突

        shardingRuleConfig.getKeyGenerators().put("snowflake",
                new ShardingSphereAlgorithmConfiguration("SNOWFLAKE", snowflakeProps));

        // 5. 创建数据源（添加配置属性）
        Properties props = new Properties();
        props.setProperty("sql-show", "true");
        props.setProperty("allow-range-query-with-inline-sharding", "true");

        return ShardingSphereDataSourceFactory.createDataSource(
                "logic_db",
                null,
                dataSourceMap,
                Collections.singleton(shardingRuleConfig),
                props);
    }

}